package com.zeratul.common.pool.utils;

import com.zeratul.common.core.utils.spring.SpringUtils;
import com.zeratul.common.pool.decorator.ZtTaskDecorator;
import com.zeratul.common.pool.decorator.ZtVirtualThreadTaskExecutor;
import lombok.AccessLevel;
import lombok.NoArgsConstructor;
import lombok.extern.slf4j.Slf4j;
import org.apache.commons.lang3.concurrent.BasicThreadFactory;
import org.springframework.scheduling.concurrent.ThreadPoolTaskExecutor;

import java.util.concurrent.*;

/**
 * 线程相关工具类.
 */
@Slf4j
@NoArgsConstructor(access = AccessLevel.PRIVATE)
public class ThreadUtils {

    /**
     * 创建线程池
     * @param corePoolSize - 核心线程数
     * @param maximumPoolSize - 最大线程数
     * @param keepAliveTimeSec - 线程的存活时间（秒）
     * @param queueSize - 队列长度
     * @param prifexName - 线程前缀名称
     * @return
     */
    public static Executor newThreadPool(int corePoolSize, int maximumPoolSize, int keepAliveTimeSec, Integer queueSize, String prifexName, boolean isVirtual) {
        //判断是否开启虚拟线程，如果开启虚拟线程，则返回虚拟线程池
        if(SpringUtils.isVirtual() && isVirtual) {
            log.debug("【{}线程池初始化】使用虚拟线程池", prifexName);
            return new ZtVirtualThreadTaskExecutor("virtual-" + prifexName);
        }

        //未开启虚拟线程，则返回自定义的线程池
        //创建线程工厂
        BasicThreadFactory.Builder builder = new BasicThreadFactory.Builder().daemon(false);
        builder.namingPattern(prifexName + "-%d");
        //自定义线程池
        ThreadPoolTaskExecutor threadPoolTaskExecutor = new ThreadPoolTaskExecutor();
        // 核心线程数 - CPU核心数的2倍（适用于IO密集型）
        threadPoolTaskExecutor.setCorePoolSize(corePoolSize);
        // 最大线程数
        threadPoolTaskExecutor.setMaxPoolSize(maximumPoolSize);
        // 线程空闲时间
        threadPoolTaskExecutor.setKeepAliveSeconds(keepAliveTimeSec);
        // 队列大小
        threadPoolTaskExecutor.setQueueCapacity(queueSize);
        // 线程工厂
        threadPoolTaskExecutor.setThreadFactory(builder.build());
        // 拒绝策略
        threadPoolTaskExecutor.setRejectedExecutionHandler(new ThreadPoolExecutor.CallerRunsPolicy());
        // 等待所有任务完成再关闭线程池
        threadPoolTaskExecutor.setWaitForTasksToCompleteOnShutdown(true);
        threadPoolTaskExecutor.setAwaitTerminationSeconds(60);
        // 设置任务装饰器
        threadPoolTaskExecutor.setTaskDecorator(new ZtTaskDecorator());
        // 初始化线程池
        threadPoolTaskExecutor.initialize();

        log.debug("【{}线程池初始化】CorePool:{} MaxPool:{} QueueSize:{}",
                prifexName,
                threadPoolTaskExecutor.getCorePoolSize(),
                threadPoolTaskExecutor.getMaxPoolSize(),
                threadPoolTaskExecutor.getQueueCapacity());

        //返回线程池对象
        return threadPoolTaskExecutor.getThreadPoolExecutor();
    }

    /**
     * 停止线程池
     * 先使用shutdown, 停止接收新任务并尝试完成所有已存在任务.
     * 如果超时, 则调用shutdownNow, 取消在workQueue中Pending的任务,并中断所有阻塞函数.
     * 如果仍然超時，則強制退出.
     * 另对在shutdown时线程本身被调用中断做了处理.
     */
    public static void shutdownAndAwaitTermination(ExecutorService pool) {
        if (pool != null && !pool.isShutdown()) {
            pool.shutdown();
            try {
                if (!pool.awaitTermination(120, TimeUnit.SECONDS)) {
                    pool.shutdownNow();
                    if (!pool.awaitTermination(120, TimeUnit.SECONDS)) {
                        log.info("Pool did not terminate");
                    }
                }
            } catch (InterruptedException ie) {
                pool.shutdownNow();
                Thread.currentThread().interrupt();
            }
        }
    }

    /**
     * 打印线程异常信息
     */
    public static void printException(Runnable r, Throwable t) {
        if (t == null && r instanceof Future<?>) {
            try {
                Future<?> future = (Future<?>) r;
                if (future.isDone()) {
                    future.get();
                }
            } catch (CancellationException ce) {
                t = ce;
            } catch (ExecutionException ee) {
                t = ee.getCause();
            } catch (InterruptedException ie) {
                Thread.currentThread().interrupt();
            }
        }
        if (t != null) {
            log.error(t.getMessage(), t);
        }
    }
}
